{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pyscf\n",
    "import numpy\n",
    "from pyscf import gto, scf, ci,df,lib, dft\n",
    "import scipy\n",
    "from time import ctime, time\n",
    "import numgrid"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {
    "scrolled": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "System: uname_result(system='Linux', node='igor-Latitude-5591', release='5.4.0-122-generic', version='#138-Ubuntu SMP Wed Jun 22 15:00:31 UTC 2022', machine='x86_64', processor='x86_64')  Threads 6\n",
      "Python 3.6.8 |Anaconda, Inc.| (default, Dec 30 2018, 01:22:34) \n",
      "[GCC 7.3.0]\n",
      "numpy 1.16.4  scipy 1.2.1\n",
      "Date: Tue Aug  9 13:17:06 2022\n",
      "PySCF version 2.0.1\n",
      "PySCF path  /home/igor/Data/opt/anaconda3/lib/python3.6/site-packages/pyscf\n",
      "\n",
      "[CONFIG] conf_file None\n",
      "[INPUT] verbose = 4\n",
      "[INPUT] num. atoms = 4\n",
      "[INPUT] num. electrons = 10\n",
      "[INPUT] charge = 0\n",
      "[INPUT] spin (= nelec alpha-beta = 2S) = 0\n",
      "[INPUT] symmetry False subgroup None\n",
      "[INPUT] Mole.unit = angstrom\n",
      "[INPUT]  1 N     -2.198839101900   1.897374626800   0.000000000000 AA   -4.155203694576   3.585518400351   0.000000000000 Bohr\n",
      "[INPUT]  2 H     -1.178839101900   1.897374626800   0.000000000000 AA   -2.227683047519   3.585518400351   0.000000000000 Bohr\n",
      "[INPUT]  3 H     -2.538835398700   1.092546014400  -0.526358644600 AA   -4.797703578894   2.064612745701  -0.994673681591 Bohr\n",
      "[INPUT]  4 H     -2.538840027600   2.755627174500  -0.433822469400 AA   -4.797712326247   5.207380661214  -0.819805653849 Bohr\n",
      "\n",
      "nuclear repulsion = 11.8479237868993\n",
      "number of shells = 44\n",
      "number of NR pGTOs = 164\n",
      "number of NR cGTOs = 145\n",
      "basis = cc-pvqz\n",
      "ecp = {}\n",
      "CPU time:        55.92\n",
      "\n",
      "\n",
      "******** <class 'pyscf.dft.rks.RKS'> ********\n",
      "method = RKS-RHF\n",
      "initial guess = 1e\n",
      "damping factor = 0\n",
      "level_shift factor = 0\n",
      "DIIS = <class 'pyscf.scf.diis.CDIIS'>\n",
      "diis_start_cycle = 1\n",
      "diis_space = 8\n",
      "SCF conv_tol = 1e-09\n",
      "SCF conv_tol_grad = None\n",
      "SCF max_cycles = 50\n",
      "direct_scf = True\n",
      "direct_scf_tol = 1e-13\n",
      "chkfile to save SCF result = /home/igor/Documents/Package-Pool/Rust/rest/example/DFT-NH3/tmp7r7g0ley\n",
      "max_memory 4000 MB (current use 323 MB)\n",
      "XC library pyscf.dft.libxc version 5.1.7\n",
      "    S. Lehtola, C. Steigemann, M. J. Oliveira, and M. A. Marques, SoftwareX 7, 1 (2018)\n",
      "XC functionals = lda, vwn\n",
      "    P. A. M. Dirac, Math. Proc. Cambridge Philos. Soc. 26, 376 (1930)\n",
      "    F. Bloch, Z. Phys. 57, 545 (1929)\n",
      "    S. H. Vosko, L. Wilk, and M. Nusair, Can. J. Phys. 58, 1200 (1980)\n",
      "radial grids: \n",
      "    Treutler-Ahlrichs [JCP 102, 346 (1995); DOI:10.1063/1.469408] (M4) radial grids\n",
      "    \n",
      "becke partition: Becke, JCP 88, 2547 (1988); DOI:10.1063/1.454033\n",
      "pruning grids: <function nwchem_prune at 0x1521125dd730>\n",
      "grids dens level: 3\n",
      "symmetrized grids: False\n",
      "atomic radii adjust function: <function treutler_atomic_radii_adjust at 0x1521125dd6a8>\n",
      "small_rho_cutoff = 1e-07\n",
      "Set gradient conv threshold to 3.16228e-05\n",
      "Initial guess from hcore.\n",
      "  HOMO = -7.54080997261236  LUMO = -4.62646850461635\n",
      "tot grids = 44118\n",
      "init E= -40.534485422395\n",
      "  HOMO = 0.251012667782133  LUMO = 0.331256645732179\n",
      "cycle= 1 E= -47.931631691598  delta_E= -7.4  |g|= 3.88  |ddm|=  204\n",
      "  HOMO = -2.66236020992194  LUMO = -1.07434393758569\n",
      "cycle= 2 E= -50.0011876371955  delta_E= -2.07  |g|= 6.46  |ddm|=  204\n",
      "  HOMO = -1.03938589528758  LUMO = -0.329771818008664\n",
      "cycle= 3 E= -55.2556691476139  delta_E= -5.25  |g|= 1.97  |ddm|= 1.77\n",
      "  HOMO = -0.557673928339623  LUMO = -0.0996842392003698\n",
      "cycle= 4 E= -55.7028937571798  delta_E= -0.447  |g|= 1.52  |ddm|= 0.486\n",
      "  HOMO = -0.315892170137481  LUMO = -0.0357402010488906\n",
      "cycle= 5 E= -56.0470542423236  delta_E= -0.344  |g|= 0.575  |ddm|= 0.486\n",
      "  HOMO = -0.234036978137258  LUMO = -0.00379819472513065\n",
      "cycle= 6 E= -56.1052810091047  delta_E= -0.0582  |g|= 0.0891  |ddm|= 0.245\n",
      "  HOMO = -0.215127452444616  LUMO = 0.00120246738914654\n",
      "cycle= 7 E= -56.1066912243103  delta_E= -0.00141  |g|= 0.0165  |ddm|= 0.121\n",
      "  HOMO = -0.21764353616461  LUMO = 0.00167960025563142\n",
      "cycle= 8 E= -56.1067558589397  delta_E= -6.46e-05  |g|= 0.00112  |ddm|= 0.0276\n",
      "  HOMO = -0.217588903269715  LUMO = 0.00172171555010488\n",
      "cycle= 9 E= -56.106756072734  delta_E= -2.14e-07  |g|= 0.000449  |ddm|= 0.000891\n",
      "  HOMO = -0.217511626670823  LUMO = 0.00174764102437496\n",
      "cycle= 10 E= -56.1067561130542  delta_E= -4.03e-08  |g|= 6.59e-06  |ddm|= 0.000323\n",
      "  HOMO = -0.217512425656009  LUMO = 0.00174832273665945\n",
      "cycle= 11 E= -56.106756113062  delta_E= -7.74e-12  |g|= 8.76e-07  |ddm|= 1.04e-05\n",
      "  HOMO = -0.217511983538389  LUMO = 0.00174848571790384\n",
      "Extra cycle  E= -56.1067561130615  delta_E= 4.97e-13  |g|= 1.68e-06  |ddm|= 1.89e-06\n",
      "converged SCF energy = -56.1067561130615\n",
      "Default DFT(LDA).  E = -56.106756113061\n",
      "Total job time:       1.41(wall)\n"
     ]
    }
   ],
   "source": [
    "TimeStart = time()\n",
    "mol = gto.Mole(                                                                                               \n",
    "        atom='''\n",
    "         N       -2.1988391019      1.8973746268      0.0000000000                 \n",
    "         H       -1.1788391019      1.8973746268      0.0000000000                 \n",
    "         H       -2.5388353987      1.0925460144     -0.5263586446                 \n",
    "         H       -2.5388400276      2.7556271745     -0.4338224694 ''',\n",
    "        basis='cc-pvqz',verbose=4\n",
    "    ).build()\n",
    "method = dft.RKS(mol)\n",
    "method.xc = 'lda, vwn'\n",
    "method.init_guess = '1e'\n",
    "method.grids.becke_scheme = dft.original_becke\n",
    "print('Default DFT(LDA).  E = %.12f' % method.kernel())\n",
    "print(\"Total job time: %10.2f(wall)\" %(time()-TimeStart))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "36396\n"
     ]
    }
   ],
   "source": [
    "# store the grids to an external file\n",
    "print(len(method.grids.weights))\n",
    "ff = open(\"grids-qz\",'w')\n",
    "for i in range(len(method.grids.weights)):\n",
    "    x,y,z = method.grids.coords[i]\n",
    "    w = method.grids.weights[i]\n",
    "    ff.write(\"%16.8e,%16.8e,%16.8e,%16.8e\\n\" %(x,y,z,w))\n",
    "ff.close()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "System: uname_result(system='Linux', node='igor-Latitude-5591', release='5.4.0-122-generic', version='#138-Ubuntu SMP Wed Jun 22 15:00:31 UTC 2022', machine='x86_64', processor='x86_64')  Threads 6\n",
      "Python 3.6.8 |Anaconda, Inc.| (default, Dec 30 2018, 01:22:34) \n",
      "[GCC 7.3.0]\n",
      "numpy 1.16.4  scipy 1.2.1\n",
      "Date: Tue Aug  9 13:17:01 2022\n",
      "PySCF version 2.0.1\n",
      "PySCF path  /home/igor/Data/opt/anaconda3/lib/python3.6/site-packages/pyscf\n",
      "\n",
      "[CONFIG] conf_file None\n",
      "[INPUT] verbose = 4\n",
      "[INPUT] num. atoms = 4\n",
      "[INPUT] num. electrons = 10\n",
      "[INPUT] charge = 0\n",
      "[INPUT] spin (= nelec alpha-beta = 2S) = 0\n",
      "[INPUT] symmetry False subgroup None\n",
      "[INPUT] Mole.unit = angstrom\n",
      "[INPUT]  1 N     -2.198839101900   1.897374626800   0.000000000000 AA   -4.155203694576   3.585518400351   0.000000000000 Bohr\n",
      "[INPUT]  2 H     -1.178839101900   1.897374626800   0.000000000000 AA   -2.227683047519   3.585518400351   0.000000000000 Bohr\n",
      "[INPUT]  3 H     -2.538835398700   1.092546014400  -0.526358644600 AA   -4.797703578894   2.064612745701  -0.994673681591 Bohr\n",
      "[INPUT]  4 H     -2.538840027600   2.755627174500  -0.433822469400 AA   -4.797712326247   5.207380661214  -0.819805653849 Bohr\n",
      "\n",
      "nuclear repulsion = 11.8479237868993\n",
      "number of shells = 86\n",
      "number of NR pGTOs = 393\n",
      "number of NR cGTOs = 367\n",
      "basis = aug-cc-pv5z\n",
      "ecp = {}\n",
      "CPU time:        37.72\n",
      "\n",
      "\n",
      "******** <class 'pyscf.df.df_jk.density_fit.<locals>.DFHF'> ********\n",
      "method = DFHF-RHF\n",
      "initial guess = 1e\n",
      "damping factor = 0\n",
      "level_shift factor = 0\n",
      "DIIS = <class 'pyscf.scf.diis.CDIIS'>\n",
      "diis_start_cycle = 1\n",
      "diis_space = 8\n",
      "SCF conv_tol = 1e-09\n",
      "SCF conv_tol_grad = None\n",
      "SCF max_cycles = 50\n",
      "direct_scf = True\n",
      "direct_scf_tol = 1e-13\n",
      "chkfile to save SCF result = /home/igor/Documents/Package-Pool/Rust/rest/example/DFT-NH3/tmpmt387o3f\n",
      "max_memory 4000 MB (current use 652 MB)\n",
      "Set gradient conv threshold to 3.16228e-05\n",
      "Initial guess from hcore.\n",
      "  HOMO = -7.55161757048187  LUMO = -4.62874347532264\n",
      "******** <class 'pyscf.df.df.DF'> ********\n",
      "auxbasis = def2-tzvp-jkfit\n",
      "max_memory = 4000\n",
      "_cderi_to_save = /home/igor/Documents/Package-Pool/Rust/rest/example/DFT-NH3/tmpznun2n_4\n",
      "init E= -40.4732793889152\n",
      "  HOMO = 0.0489561315679074  LUMO = 0.0830079472308457\n",
      "cycle= 1 E= -44.3756896713036  delta_E= -3.9  |g|= 3.51  |ddm|=  457\n",
      "  HOMO = -3.27597406570694  LUMO = -1.79113256766857\n",
      "cycle= 2 E= -51.3664875422616  delta_E= -6.99  |g|=  5.4  |ddm|=  457\n",
      "  HOMO = -1.55582223080992  LUMO = -0.723940590706147\n",
      "cycle= 3 E= -55.1758195011046  delta_E= -3.81  |g|= 2.18  |ddm|= 1.33\n",
      "  HOMO = -0.861171177842922  LUMO = -0.122616028714864\n",
      "cycle= 4 E= -55.9782147176332  delta_E= -0.802  |g|= 1.03  |ddm|= 0.559\n",
      "  HOMO = -0.445564677855217  LUMO = 0.0135015839441978\n",
      "cycle= 5 E= -56.2222386073672  delta_E= -0.244  |g|= 0.053  |ddm|= 0.431\n",
      "  HOMO = -0.415660016354948  LUMO = 0.024163811694455\n",
      "cycle= 6 E= -56.2237209625422  delta_E= -0.00148  |g|= 0.0172  |ddm|= 0.093\n",
      "  HOMO = -0.419692553365013  LUMO = 0.0239824185508774\n",
      "cycle= 7 E= -56.2238200686541  delta_E= -9.91e-05  |g|= 0.00532  |ddm|= 0.0256\n",
      "  HOMO = -0.421543491210618  LUMO = 0.0235299751279285\n",
      "cycle= 8 E= -56.2238306074276  delta_E= -1.05e-05  |g|= 0.00071  |ddm|= 0.00557\n",
      "  HOMO = -0.421242532536969  LUMO = 0.023528403202773\n",
      "cycle= 9 E= -56.2238309108378  delta_E= -3.03e-07  |g|= 0.000271  |ddm|= 0.00116\n",
      "  HOMO = -0.421107836794169  LUMO = 0.0235562691562148\n",
      "cycle= 10 E= -56.2238309346254  delta_E= -2.38e-08  |g|= 3.29e-05  |ddm|= 0.000432\n",
      "  HOMO = -0.421101635359577  LUMO = 0.023556419485649\n",
      "cycle= 11 E= -56.2238309352643  delta_E= -6.39e-10  |g|= 6.31e-06  |ddm|= 7.87e-05\n",
      "  HOMO = -0.421103963595954  LUMO = 0.023556417629366\n",
      "Extra cycle  E= -56.2238309352784  delta_E= -1.41e-11  |g|= 3.66e-06  |ddm|= 1.11e-05\n",
      "converged SCF energy = -56.2238309352784\n",
      "Total energy: -56.22383093527842\n",
      "Total job time:       3.29(wall)\n"
     ]
    }
   ],
   "source": [
    "lib.num_threads(6)\n",
    "TimeStart = time()\n",
    "mol = gto.Mole(                                                                                               \n",
    "        atom='''\n",
    "         N       -2.1988391019      1.8973746268      0.0000000000                 \n",
    "         H       -1.1788391019      1.8973746268      0.0000000000                 \n",
    "         H       -2.5388353987      1.0925460144     -0.5263586446                 \n",
    "         H       -2.5388400276      2.7556271745     -0.4338224694 ''',\n",
    "        charge=0,spin=0,basis='aug-cc-pv5z',verbose=4\n",
    "    ).build()\n",
    "mf = scf.RHF(mol).density_fit(auxbasis=\"def2-tzvp-jkfit\")\n",
    "mf.init_guess = '1e'\n",
    "mf.diis = 'diis'\n",
    "print(\"Total energy:\",mf.kernel())\n",
    "print(\"Total job time: %10.2f(wall)\" %(time()-TimeStart))\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "ename": "NameError",
     "evalue": "name 'mf' is not defined",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mNameError\u001b[0m                                 Traceback (most recent call last)",
      "\u001b[0;32m<ipython-input-7-eb8637577f1d>\u001b[0m in \u001b[0;36m<module>\u001b[0;34m\u001b[0m\n\u001b[1;32m      1\u001b[0m \u001b[0;31m# store the grids to an external file\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m----> 2\u001b[0;31m \u001b[0mprint\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mlen\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mmf\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mgrids\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mweights\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m      3\u001b[0m \u001b[0mff\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mopen\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m\"grids-qz\"\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m'w'\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m      4\u001b[0m \u001b[0;32mfor\u001b[0m \u001b[0mi\u001b[0m \u001b[0;32min\u001b[0m \u001b[0mrange\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mlen\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mmethod\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mgrids\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mweights\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m      5\u001b[0m     \u001b[0mx\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0my\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0mz\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mmethod\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mgrids\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mcoords\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0mi\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;31mNameError\u001b[0m: name 'mf' is not defined"
     ]
    }
   ],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Default DFT(LDA).  E = -56.100667658113\n",
      "Changed radial grids for DFT.  E = -56.100667683944\n",
      "Changed grid partition funciton.  E = -56.100667820637\n",
      "Dense grids.  E = -56.100667719163\n",
      "Dense grids for O atom.  E = -56.100667658113\n",
      "Dense grids for all atoms.  E = -56.100667729458\n",
      "Changed grid partition funciton.  E = -56.100667639101\n",
      "Total job time:       2.85(wall)\n"
     ]
    }
   ],
   "source": [
    "TimeStart = time()\n",
    "mol = gto.Mole(                                                                                               \n",
    "        atom='''\n",
    "         N       -2.1988391019      1.8973746268      0.0000000000                 \n",
    "         H       -1.1788391019      1.8973746268      0.0000000000                 \n",
    "         H       -2.5388353987      1.0925460144     -0.5263586446                 \n",
    "         H       -2.5388400276      2.7556271745     -0.4338224694 ''',\n",
    "        basis='def2-tzvp',verbose=0\n",
    "    ).build()\n",
    "method = dft.RKS(mol)\n",
    "print('Default DFT(LDA).  E = %.12f' % method.kernel())\n",
    "\n",
    "# See pyscf/dft/radi.py for more radial grid schemes\n",
    "#grids.radi_method = dft.gauss_chebeshev\n",
    "#grids.radi_method = dft.delley\n",
    "method = dft.RKS(mol)\n",
    "method.grids.radi_method = dft.mura_knowles\n",
    "print('Changed radial grids for DFT.  E = %.12f' % method.kernel())\n",
    "\n",
    "\n",
    "# See pyscf/dft/gen_grid.py for detail of the grid weight scheme\n",
    "#method.grids.becke_scheme = dft.original_becke\n",
    "# Stratmann-Scuseria weight scheme\n",
    "method = dft.RKS(mol)\n",
    "method.grids.becke_scheme = dft.stratmann\n",
    "print('Changed grid partition funciton.  E = %.12f' % method.kernel())\n",
    "\n",
    "# Grids level 0 - 9.  Big number indicates dense grids. Default is 3\n",
    "method = dft.RKS(mol)\n",
    "method.grids.level = 4\n",
    "print('Dense grids.  E = %.12f' % method.kernel())\n",
    "\n",
    "# Specify mesh grid for certain atom\n",
    "method = dft.RKS(mol)\n",
    "method.grids.atom_grid = {'O': (100, 770)}\n",
    "print('Dense grids for O atom.  E = %.12f' % method.kernel())\n",
    "\n",
    "# Specify mesh grid for all atoms\n",
    "method = dft.RKS(mol)\n",
    "method.grids.atom_grid = (100, 770)\n",
    "print('Dense grids for all atoms.  E = %.12f' % method.kernel())\n",
    "\n",
    "\n",
    "# Disable pruning grids near core region\n",
    "#grids.prune = dft.sg1_prune\n",
    "method = dft.RKS(mol)\n",
    "method.grids.prune = None\n",
    "print('Changed grid partition funciton.  E = %.12f' % method.kernel())\n",
    "\n",
    "print(\"Total job time: %10.2f(wall)\" %(time()-TimeStart))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([3.77945036, 0.66140414, 2.64561657, 2.74010288, 1.98421243,\n",
       "       1.60626721, 1.32280829, 1.22832198, 1.13383567, 0.94486306,\n",
       "       2.83458919, 3.40150702, 2.83458919, 2.36215766, 2.07869874,\n",
       "       1.88972612, 1.88972612, 1.88972612, 3.40150702, 4.15739747,\n",
       "       3.40150702, 3.0235618 , 2.64561657, 2.55113027, 2.64561657,\n",
       "       2.64561657, 2.64561657, 2.55113027, 2.55113027, 2.55113027,\n",
       "       2.55113027, 2.45664396, 2.36215766, 2.17318504, 2.17318504,\n",
       "       2.17318504, 3.59047964, 4.44085639, 3.77945225, 3.40150702,\n",
       "       2.92907549, 2.74010288, 2.74010288, 2.55113027, 2.45664396,\n",
       "       2.55113027, 2.64561657, 3.0235618 , 2.92907549, 2.92907549,\n",
       "       2.74010288, 2.74010288, 2.64561657, 2.64561657, 3.96842486,\n",
       "       4.91328792, 4.06291117, 3.68496594, 3.49599333, 3.49599333,\n",
       "       3.49599333, 3.49599333, 3.49599333, 3.49599333, 3.40150702,\n",
       "       3.30702072, 3.30702072, 3.30702072, 3.30702072, 3.30702072,\n",
       "       3.30702072, 3.30702072, 2.92907549, 2.74010288, 2.55113027,\n",
       "       2.55113027, 2.45664396, 2.55113027, 2.55113027, 2.55113027,\n",
       "       2.83458919, 3.59047964, 3.40150702, 3.0235618 , 3.59047964,\n",
       "       2.74010288, 3.96842486, 3.40150702, 4.06291117, 3.68496594,\n",
       "       3.40150702, 3.40150702, 3.30702072, 3.30702072, 3.30702072,\n",
       "       3.30702072, 3.30702072, 3.30702072, 3.30702072, 3.30702072,\n",
       "       3.30702072, 3.30702072, 3.30702072, 3.30702072, 3.30702072,\n",
       "       3.30702072, 3.30702072, 3.30702072, 3.30702072, 3.30702072,\n",
       "       3.30702072, 3.30702072, 3.30702072, 3.30702072, 3.30702072,\n",
       "       3.30702072, 3.30702072, 3.30702072, 3.30702072, 3.30702072,\n",
       "       3.30702072, 3.30702072, 3.30702072, 3.30702072, 3.30702072,\n",
       "       3.30702072, 3.30702072, 3.30702072, 3.30702072, 3.30702072,\n",
       "       3.30702072])"
      ]
     },
     "execution_count": 26,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "grids = dft.gen_grid.Grids(mol)\n",
    "grids.atomic_radii\n",
    "#help(dft.gen_grid.Grids)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'N': (array([[ 8.88723093e-06,  0.00000000e+00,  0.00000000e+00],\n",
       "         [ 8.16664970e-05,  0.00000000e+00,  0.00000000e+00],\n",
       "         [ 2.98885715e-04,  0.00000000e+00,  0.00000000e+00],\n",
       "         ...,\n",
       "         [-3.23900501e+00, -1.48150365e+00, -4.42739538e-01],\n",
       "         [-3.44745973e+00, -1.57684973e+00, -4.71233209e-01],\n",
       "         [-3.66837117e+00, -1.67789344e+00, -5.01429590e-01]]),\n",
       "  array([3.58430380e-16, 1.39057893e-13, 4.54432003e-12, ...,\n",
       "         1.23267854e-01, 1.47869760e-01, 1.77583517e-01])),\n",
       " 'H': (array([[ 3.18520842e-05,  0.00000000e+00,  0.00000000e+00],\n",
       "         [ 2.92679823e-04,  0.00000000e+00,  0.00000000e+00],\n",
       "         [ 1.07106729e-03,  0.00000000e+00,  0.00000000e+00],\n",
       "         ...,\n",
       "         [-2.53740746e+00,  1.16059666e+00, -3.46838181e-01],\n",
       "         [ 2.53740746e+00, -1.16059666e+00, -3.46838181e-01],\n",
       "         [-2.53740746e+00, -1.16059666e+00, -3.46838181e-01]]),\n",
       "  array([1.65011524e-14, 6.40064128e-12, 2.09103831e-10, ...,\n",
       "         9.14227644e-02, 9.14227644e-02, 9.14227644e-02]))}"
      ]
     },
     "execution_count": 31,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "grids.gen_atomic_grids(mol)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {
    "scrolled": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(array([3.18520842e-05, 2.92679823e-04, 1.07106729e-03, 2.68860020e-03,\n",
      "       5.48931493e-03, 9.83452575e-03, 1.60996016e-02, 2.46717290e-02,\n",
      "       3.59482665e-02, 5.03355017e-02, 6.82477101e-02, 9.01064631e-02,\n",
      "       1.16340157e-01, 1.47383751e-01, 1.83678726e-01, 2.25673256e-01,\n",
      "       2.73822632e-01, 3.28589962e-01, 3.90447176e-01, 4.59876402e-01,\n",
      "       5.37371759e-01, 6.23441661e-01, 7.18611726e-01, 8.23428409e-01,\n",
      "       9.38463534e-01, 1.06431991e+00, 1.20163833e+00, 1.35110624e+00,\n",
      "       1.51346858e+00, 1.68954138e+00, 1.88022891e+00, 2.08654545e+00,\n",
      "       2.30964331e+00, 2.55084900e+00, 2.81171086e+00, 3.09406227e+00,\n",
      "       3.40010722e+00, 3.73253801e+00, 4.09470078e+00, 4.49083388e+00,\n",
      "       4.92642110e+00, 5.40873286e+00, 5.94768924e+00, 6.55730490e+00,\n",
      "       7.25825956e+00, 8.08284042e+00, 9.08548337e+00, 1.03688003e+01,\n",
      "       1.21642006e+01, 1.52126401e+01]), array([1.01924658e-04, 4.68250899e-04, 1.14227106e-03, 2.15021492e-03,\n",
      "       3.51148541e-03, 5.24153418e-03, 7.35317569e-03, 9.85732603e-03,\n",
      "       1.27634791e-02, 1.60800486e-02, 1.98146379e-02, 2.39742716e-02,\n",
      "       2.85656096e-02, 3.35951569e-02, 3.90694817e-02, 4.49954476e-02,\n",
      "       5.13804725e-02, 5.82328197e-02, 6.55619340e-02, 7.33788327e-02,\n",
      "       8.16965683e-02, 9.05307791e-02, 9.99003518e-02, 1.09828225e-01,\n",
      "       1.20342377e-01, 1.31477036e-01, 1.43274201e-01, 1.55785548e-01,\n",
      "       1.69074851e-01, 1.83221098e-01, 1.98322538e-01, 2.14502003e-01,\n",
      "       2.31914014e-01, 2.50754391e-01, 2.71273489e-01, 2.93794741e-01,\n",
      "       3.18741174e-01, 3.46674219e-01, 3.78351986e-01, 4.14819450e-01,\n",
      "       4.57552908e-01, 5.08701042e-01, 5.71507234e-01, 6.51094708e-01,\n",
      "       7.56038309e-01, 9.01825468e-01, 1.11951530e+00, 1.48173342e+00,\n",
      "       2.20671010e+00, 4.38688533e+00]))\n",
      "[2.1683888135266817e-07, 6.901344678372829e-06, 5.193644169576905e-05, 0.00021612001435303315, 0.0006489967020554431, 0.001583609760623531, 0.003345267980792065, 0.006354043094066759, 0.011121358804858978, 0.018241602973560568, 0.028380103841055576, 0.042259010701481824, 0.060642606265369145, 0.0843233899331677, 0.1141099711116606, 0.15081746836463844, 0.1952607827134402, 0.24825084231985833, 0.31059372121657636, 0.38309241972605645, 0.4665510503947239, 0.5617811868686836, 0.6696101896088374, 0.7908914095071005, 0.9265162801675565, 1.0774284384148651, 1.2446401619859342, 1.42925159003529, 1.6324734085911998, 1.8556539593879195, 2.1003120968395663, 2.368177620057018, 2.661241814797698, 2.981821663600723, 3.3326427964439036, 3.716948547162165, 4.13864603938759, 4.602505898094958, 5.114441498157418, 5.6819094848355, 6.31450124252428, 7.02484660301932, 7.830051638454013, 8.754101182318156, 9.832125525459016, 11.11859209164288, 12.704754128261852, 14.76168564101178, 17.673460904655776, 22.665639676246034]\n",
      "[5.094100298064615e-20, 8.193786101878449e-16, 2.319732703903635e-13, 1.247306108215661e-11, 2.6847203603240513e-10, 3.2250577507193813e-09, 2.582154736987662e-08, 1.5324076026145257e-07, 7.221149276836952e-07, 2.8328386193423405e-06, 9.572014850053967e-06, 2.8575111446471595e-05, 7.685439692484097e-05, 0.0001891240233338212, 0.0004311415486847432, 0.0009198410017747929, 0.0018522731280789823, 0.00354563939698269, 0.006491077084431341, 0.011424409613217595, 0.019417919898872905, 0.03199845412961261, 0.051298971146082666, 0.08025322683276731, 0.12284692554565575, 0.18444383239588236, 0.27221270626992655, 0.3956915381022138, 0.5675411051326515, 0.8045628769076054, 1.129091006308462, 1.5709213418235999, 2.17002352606663, 2.9804148300587876, 4.075790664966723, 5.557868487142548, 7.569023995843895, 10.311902491337063, 14.080717453128653, 19.312831451292922, 26.676997878273045, 37.2311103406916, 52.719452531480464, 76.1699333489475, 113.19414812025786, 175.1249536545187, 287.71507965283735, 520.5470590005915, 1123.3114228013753, 3703.1089028101605]\n"
     ]
    }
   ],
   "source": [
    "grids = dft.gen_grid.Grids(mol)\n",
    "print(grids.radi_method(50))\n",
    "#radii, weights = numgrid.radial_grid_lmg_bse( basis_set=\"sto-3g\", radial_precision=1.0e-12, proton_charge=1)\n",
    "radii, weights = numgrid.radial_grid_kk(num_points=50)\n",
    "print(radii)\n",
    "print(weights)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "scrolled": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Help on Grids in module pyscf.dft.gen_grid object:\n",
      "\n",
      "class Grids(pyscf.lib.misc.StreamObject)\n",
      " |  DFT mesh grids\n",
      " |  \n",
      " |  Attributes for Grids:\n",
      " |      level : int\n",
      " |          To control the number of radial and angular grids. Large number\n",
      " |          leads to large mesh grids. The default level 3 corresponds to\n",
      " |          (50,302) for H, He;\n",
      " |          (75,302) for second row;\n",
      " |          (80~105,434) for rest.\n",
      " |  \n",
      " |          Grids settings at other levels can be found in\n",
      " |          pyscf.dft.gen_grid.RAD_GRIDS and pyscf.dft.gen_grid.ANG_ORDER\n",
      " |  \n",
      " |      atomic_radii : 1D array\n",
      " |          | radi.BRAGG_RADII  (default)\n",
      " |          | radi.COVALENT_RADII\n",
      " |          | None : to switch off atomic radii adjustment\n",
      " |  \n",
      " |      radii_adjust : function(mol, atomic_radii) => (function(atom_id, atom_id, g) => array_like_g)\n",
      " |          Function to adjust atomic radii, can be one of\n",
      " |          | radi.treutler_atomic_radii_adjust\n",
      " |          | radi.becke_atomic_radii_adjust\n",
      " |          | None : to switch off atomic radii adjustment\n",
      " |  \n",
      " |      radi_method : function(n) => (rad_grids, rad_weights)\n",
      " |          scheme for radial grids, can be one of\n",
      " |          | radi.treutler  (default)\n",
      " |          | radi.delley\n",
      " |          | radi.mura_knowles\n",
      " |          | radi.gauss_chebyshev\n",
      " |  \n",
      " |      becke_scheme : function(v) => array_like_v\n",
      " |          weight partition function, can be one of\n",
      " |          | gen_grid.original_becke  (default)\n",
      " |          | gen_grid.stratmann\n",
      " |  \n",
      " |      prune : function(nuc, rad_grids, n_ang) => list_n_ang_for_each_rad_grid\n",
      " |          scheme to reduce number of grids, can be one of\n",
      " |          | gen_grid.nwchem_prune  (default)\n",
      " |          | gen_grid.sg1_prune\n",
      " |          | gen_grid.treutler_prune\n",
      " |          | None : to switch off grid pruning\n",
      " |  \n",
      " |      symmetry : bool\n",
      " |          whether to symmetrize mesh grids (TODO)\n",
      " |  \n",
      " |      atom_grid : dict\n",
      " |          Set (radial, angular) grids for particular atoms.\n",
      " |          Eg, grids.atom_grid = {'H': (20,110)} will generate 20 radial\n",
      " |          grids and 110 angular grids for H atom.\n",
      " |  \n",
      " |      Examples:\n",
      " |  \n",
      " |      >>> mol = gto.M(atom='H 0 0 0; H 0 0 1.1')\n",
      " |      >>> grids = dft.gen_grid.Grids(mol)\n",
      " |      >>> grids.level = 4\n",
      " |      >>> grids.build()\n",
      " |  \n",
      " |  Method resolution order:\n",
      " |      Grids\n",
      " |      pyscf.lib.misc.StreamObject\n",
      " |      builtins.object\n",
      " |  \n",
      " |  Methods defined here:\n",
      " |  \n",
      " |  __init__(self, mol)\n",
      " |      Initialize self.  See help(type(self)) for accurate signature.\n",
      " |  \n",
      " |  __setattr__(self, key, val)\n",
      " |      Implement setattr(self, name, value).\n",
      " |  \n",
      " |  build(self, mol=None, with_non0tab=False, **kwargs)\n",
      " |  \n",
      " |  dump_flags(self, verbose=None)\n",
      " |  \n",
      " |  gen_atomic_grids(self, mol, atom_grid=None, radi_method=None, level=None, prune=None, **kwargs)\n",
      " |      Generate number of radial grids and angular grids for the given molecule.\n",
      " |      \n",
      " |      Returns:\n",
      " |          A dict, with the atom symbol for the dict key.  For each atom type,\n",
      " |          the dict value has two items: one is the meshgrid coordinates wrt the\n",
      " |          atom center; the second is the volume of that grid.\n",
      " |  \n",
      " |  gen_partition = get_partition(self, mol, atom_grids_tab=None, radii_adjust=None, atomic_radii=array([3.77945036, 0.66140414, 2.64561657, 2.74010288, 1.98421243,\n",
      " |         1.60626721, 1.32280829, 1.22832198, 1.13383567, 0.94486306,\n",
      " |         2.83458919, 3.40150702, 2.83458919, 2.36215766, 2.07869874,\n",
      " |         1.88972612, 1.88972612, 1.88972612, 3.40150702, 4.15739747,\n",
      " |         3.40150702, 3.0235618 , 2.64561657, 2.55113027, 2.64561657,\n",
      " |         2.64561657, 2.64561657, 2.55113027, 2.55113027, 2.55113027,\n",
      " |         2.55113027, 2.45664396, 2.36215766, 2.17318504, 2.17318504,\n",
      " |         2.17318504, 3.59047964, 4.44085639, 3.77945225, 3.40150702,\n",
      " |         2.92907549, 2.74010288, 2.74010288, 2.55113027, 2.45664396,\n",
      " |         2.55113027, 2.64561657, 3.0235618 , 2.92907549, 2.92907549,\n",
      " |         2.74010288, 2.74010288, 2.64561657, 2.64561657, 3.96842486,\n",
      " |         4.91328792, 4.06291117, 3.68496594, 3.49599333, 3.49599333,\n",
      " |         3.49599333, 3.49599333, 3.49599333, 3.49599333, 3.40150702,\n",
      " |         3.30702072, 3.30702072, 3.30702072, 3.30702072, 3.30702072,\n",
      " |         3.30702072, 3.30702072, 2.92907549, 2.74010288, 2.55113027,\n",
      " |         2.55113027, 2.45664396, 2.55113027, 2.55113027, 2.55113027,\n",
      " |         2.83458919, 3.59047964, 3.40150702, 3.0235618 , 3.59047964,\n",
      " |         2.74010288, 3.96842486, 3.40150702, 4.06291117, 3.68496594,\n",
      " |         3.40150702, 3.40150702, 3.30702072, 3.30702072, 3.30702072,\n",
      " |         3.30702072, 3.30702072, 3.30702072, 3.30702072, 3.30702072,\n",
      " |         3.30702072, 3.30702072, 3.30702072, 3.30702072, 3.30702072,\n",
      " |         3.30702072, 3.30702072, 3.30702072, 3.30702072, 3.30702072,\n",
      " |         3.30702072, 3.30702072, 3.30702072, 3.30702072, 3.30702072,\n",
      " |         3.30702072, 3.30702072, 3.30702072, 3.30702072, 3.30702072,\n",
      " |         3.30702072, 3.30702072, 3.30702072, 3.30702072, 3.30702072,\n",
      " |         3.30702072, 3.30702072, 3.30702072, 3.30702072, 3.30702072,\n",
      " |         3.30702072]), becke_scheme=<function original_becke at 0x1467d3af8ea0>, concat=True)\n",
      " |  \n",
      " |  get_partition(self, mol, atom_grids_tab=None, radii_adjust=None, atomic_radii=array([3.77945036, 0.66140414, 2.64561657, 2.74010288, 1.98421243,\n",
      " |         1.60626721, 1.32280829, 1.22832198, 1.13383567, 0.94486306,\n",
      " |         2.83458919, 3.40150702, 2.83458919, 2.36215766, 2.07869874,\n",
      " |         1.88972612, 1.88972612, 1.88972612, 3.40150702, 4.15739747,\n",
      " |         3.40150702, 3.0235618 , 2.64561657, 2.55113027, 2.64561657,\n",
      " |         2.64561657, 2.64561657, 2.55113027, 2.55113027, 2.55113027,\n",
      " |         2.55113027, 2.45664396, 2.36215766, 2.17318504, 2.17318504,\n",
      " |         2.17318504, 3.59047964, 4.44085639, 3.77945225, 3.40150702,\n",
      " |         2.92907549, 2.74010288, 2.74010288, 2.55113027, 2.45664396,\n",
      " |         2.55113027, 2.64561657, 3.0235618 , 2.92907549, 2.92907549,\n",
      " |         2.74010288, 2.74010288, 2.64561657, 2.64561657, 3.96842486,\n",
      " |         4.91328792, 4.06291117, 3.68496594, 3.49599333, 3.49599333,\n",
      " |         3.49599333, 3.49599333, 3.49599333, 3.49599333, 3.40150702,\n",
      " |         3.30702072, 3.30702072, 3.30702072, 3.30702072, 3.30702072,\n",
      " |         3.30702072, 3.30702072, 2.92907549, 2.74010288, 2.55113027,\n",
      " |         2.55113027, 2.45664396, 2.55113027, 2.55113027, 2.55113027,\n",
      " |         2.83458919, 3.59047964, 3.40150702, 3.0235618 , 3.59047964,\n",
      " |         2.74010288, 3.96842486, 3.40150702, 4.06291117, 3.68496594,\n",
      " |         3.40150702, 3.40150702, 3.30702072, 3.30702072, 3.30702072,\n",
      " |         3.30702072, 3.30702072, 3.30702072, 3.30702072, 3.30702072,\n",
      " |         3.30702072, 3.30702072, 3.30702072, 3.30702072, 3.30702072,\n",
      " |         3.30702072, 3.30702072, 3.30702072, 3.30702072, 3.30702072,\n",
      " |         3.30702072, 3.30702072, 3.30702072, 3.30702072, 3.30702072,\n",
      " |         3.30702072, 3.30702072, 3.30702072, 3.30702072, 3.30702072,\n",
      " |         3.30702072, 3.30702072, 3.30702072, 3.30702072, 3.30702072,\n",
      " |         3.30702072, 3.30702072, 3.30702072, 3.30702072, 3.30702072,\n",
      " |         3.30702072]), becke_scheme=<function original_becke at 0x1467d3af8ea0>, concat=True)\n",
      " |      Generate the mesh grid coordinates and weights for DFT numerical integration.\n",
      " |      We can change radii_adjust, becke_scheme functions to generate different meshgrid.\n",
      " |      \n",
      " |      Kwargs:\n",
      " |          concat: bool\n",
      " |              Whether to concatenate grids and weights in return\n",
      " |      \n",
      " |      Returns:\n",
      " |          grid_coord and grid_weight arrays.  grid_coord array has shape (N,3);\n",
      " |          weight 1D array has N elements.\n",
      " |  \n",
      " |  kernel(self, mol=None, with_non0tab=False)\n",
      " |      Kernel function is the main driver of a method.  Every method should\n",
      " |      define the kernel function as the entry of the calculation.  Note the\n",
      " |      return value of kernel function is not strictly defined.  It can be\n",
      " |      anything related to the method (such as the energy, the wave-function,\n",
      " |      the DFT mesh grids etc.).\n",
      " |  \n",
      " |  make_mask(self, mol=None, coords=None, relativity=0, shls_slice=None, verbose=None)\n",
      " |      Mask to indicate whether a shell is zero on grid\n",
      " |      \n",
      " |      Args:\n",
      " |          mol : an instance of :class:`Mole`\n",
      " |      \n",
      " |          coords : 2D array, shape (N,3)\n",
      " |              The coordinates of grids.\n",
      " |      \n",
      " |      Kwargs:\n",
      " |          relativity : bool\n",
      " |              No effects.\n",
      " |          shls_slice : 2-element list\n",
      " |              (shl_start, shl_end).\n",
      " |              If given, only part of AOs (shl_start <= shell_id < shl_end) are\n",
      " |              evaluated.  By default, all shells defined in mol will be evaluated.\n",
      " |          verbose : int or object of :class:`Logger`\n",
      " |              No effects.\n",
      " |      \n",
      " |      Returns:\n",
      " |          2D mask array of shape (N,nbas), where N is the number of grids, nbas\n",
      " |          is the number of shells.\n",
      " |  \n",
      " |  reset(self, mol=None)\n",
      " |      Reset mol and clean up relevant attributes for scanner mode\n",
      " |  \n",
      " |  ----------------------------------------------------------------------\n",
      " |  Data descriptors defined here:\n",
      " |  \n",
      " |  size\n",
      " |  \n",
      " |  ----------------------------------------------------------------------\n",
      " |  Methods inherited from pyscf.lib.misc.StreamObject:\n",
      " |  \n",
      " |  __call__ = set(self, *args, **kwargs)\n",
      " |  \n",
      " |  apply(self, fn, *args, **kwargs)\n",
      " |      Apply the fn to rest arguments:  return fn(*args, **kwargs).  The\n",
      " |      return value of method set is the object itself.  This allows a series\n",
      " |      of functions/methods to be executed in pipe.\n",
      " |  \n",
      " |  check_sanity(self)\n",
      " |      Check input of class/object attributes, check whether a class method is\n",
      " |      overwritten.  It does not check the attributes which are prefixed with\n",
      " |      \"_\".  The\n",
      " |      return value of method set is the object itself.  This allows a series\n",
      " |      of functions/methods to be executed in pipe.\n",
      " |  \n",
      " |  post_kernel(self, envs)\n",
      " |      A hook to be run after the main body of the kernel function.  Internal\n",
      " |      variables are exposed to post_kernel through the \"envs\" dictionary.\n",
      " |      Return value of post_kernel function is not required.\n",
      " |  \n",
      " |  pre_kernel(self, envs)\n",
      " |      A hook to be run before the main body of kernel function is executed.\n",
      " |      Internal variables are exposed to pre_kernel through the \"envs\"\n",
      " |      dictionary.  Return value of pre_kernel function is not required.\n",
      " |  \n",
      " |  run(self, *args, **kwargs)\n",
      " |      Call the kernel function of current object.  `args` will be passed\n",
      " |      to kernel function.  `kwargs` will be used to update the attributes of\n",
      " |      current object.  The return value of method run is the object itself.\n",
      " |      This allows a series of functions/methods to be executed in pipe.\n",
      " |  \n",
      " |  set(self, *args, **kwargs)\n",
      " |      Update the attributes of the current object.  The return value of\n",
      " |      method set is the object itself.  This allows a series of\n",
      " |      functions/methods to be executed in pipe.\n",
      " |  \n",
      " |  view(self, cls)\n",
      " |      New view of object with the same attributes.\n",
      " |  \n",
      " |  ----------------------------------------------------------------------\n",
      " |  Data descriptors inherited from pyscf.lib.misc.StreamObject:\n",
      " |  \n",
      " |  __dict__\n",
      " |      dictionary for instance variables (if defined)\n",
      " |  \n",
      " |  __weakref__\n",
      " |      list of weak references to the object (if defined)\n",
      " |  \n",
      " |  ----------------------------------------------------------------------\n",
      " |  Data and other attributes inherited from pyscf.lib.misc.StreamObject:\n",
      " |  \n",
      " |  stdout = <ipykernel.iostream.OutStream object>\n",
      " |  \n",
      " |  verbose = 0\n",
      "\n"
     ]
    }
   ],
   "source": [
    "help(grids)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 141,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "True\n",
      "True\n",
      "4.883871085326064e-13\n",
      "-2.7755575615628914e-15\n",
      "-1.10015376487841 -1.1001547208961515 -1.100193764624759 -1.1001937646247586\n"
     ]
    }
   ],
   "source": [
    "    from pyscf import scf\n",
    "    from pyscf import ao2mo\n",
    "    mol = gto.Mole()\n",
    "    mol.verbose = 0\n",
    "    mol.output = None\n",
    "\n",
    "    mol.atom.extend([\n",
    "        [\"H\", (0,  0, 0  )],\n",
    "        [\"H\", (0,  0, 1  )],\n",
    "    ])\n",
    "    mol.basis = 'cc-pvdz'\n",
    "    mol.build()\n",
    "\n",
    "    auxmol = df.incore.format_aux_basis(mol)\n",
    "    j3c = df.incore.aux_e2(mol, auxmol, intor='int3c2e_sph', aosym='s1')\n",
    "    nao = mol.nao_nr()\n",
    "    naoaux = auxmol.nao_nr()\n",
    "    j3c = j3c.reshape(nao,nao,naoaux)\n",
    "\n",
    "    atm, bas, env = \\\n",
    "            gto.mole.conc_env(mol._atm, mol._bas, mol._env,\n",
    "                              auxmol._atm, auxmol._bas, auxmol._env)\n",
    "    eri0 = numpy.empty((nao,nao,naoaux))\n",
    "    pi = 0\n",
    "    for i in range(mol.nbas):\n",
    "        pj = 0\n",
    "        for j in range(mol.nbas):\n",
    "            pk = 0\n",
    "            for k in range(mol.nbas, mol.nbas+auxmol.nbas):\n",
    "                shls = (i, j, k)\n",
    "                buf = gto.moleintor.getints_by_shell('int3c2e_sph',\n",
    "                                                     shls, atm, bas, env)\n",
    "                di, dj, dk = buf.shape\n",
    "                eri0[pi:pi+di,pj:pj+dj,pk:pk+dk] = buf\n",
    "                pk += dk\n",
    "            pj += dj\n",
    "        pi += di\n",
    "    print(numpy.allclose(eri0, j3c))\n",
    "\n",
    "    j2c = df.incore.fill_2c2e(mol, auxmol)\n",
    "    eri0 = numpy.empty_like(j2c)\n",
    "    pi = 0\n",
    "    for i in range(mol.nbas, len(bas)):\n",
    "        pj = 0\n",
    "        for j in range(mol.nbas, len(bas)):\n",
    "            shls = (i, j)\n",
    "            buf = gto.moleintor.getints_by_shell('int2c2e_sph',\n",
    "                                                 shls, atm, bas, env)\n",
    "            di, dj = buf.shape\n",
    "            eri0[pi:pi+di,pj:pj+dj] = buf\n",
    "            pj += dj\n",
    "        pi += di\n",
    "    print(numpy.allclose(eri0, j2c))\n",
    "\n",
    "    j3c = df.incore.aux_e2(mol, auxmol, intor='int3c2e_sph', aosym='s2ij')\n",
    "    cderi = df.incore.cholesky_eri(mol, auxmol=auxmol)\n",
    "    eri0 = numpy.einsum('pi,pk->ik', cderi, cderi)\n",
    "    eri1 = numpy.einsum('ik,kl->il', j3c, numpy.linalg.inv(j2c))\n",
    "    eri1 = numpy.einsum('ip,kp->ik', eri1, j3c)\n",
    "    print(abs(eri1 - eri0).max())\n",
    "    eri0 = ao2mo.restore(1, eri0, nao)\n",
    "\n",
    "    mf = scf.RHF(mol)\n",
    "    ehf0 = mf.scf()\n",
    "\n",
    "    nao = mf.mo_energy.size\n",
    "    eri1 = ao2mo.restore(1, mf._eri, nao)\n",
    "    print(abs(eri1-eri0).max() - 0.0022142583265513105)\n",
    "\n",
    "    mf._eri = ao2mo.restore(8, eri0, nao)\n",
    "    ehf1 = mf.scf()\n",
    "\n",
    "    mf = scf.RHF(mol).density_fit(auxbasis='weigend')\n",
    "    ehf2 = mf.scf()\n",
    "\n",
    "    mf = mf.density_fit(auxbasis='weigend')\n",
    "    ehf3 = mf.scf()\n",
    "    print(ehf0, ehf1, ehf2, ehf3)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[1.2304794589759454e-06, 2.6336980993310114e-06, 4.233905675553877e-06, 6.058756210858283e-06, 8.13978590080707e-06, 1.0512958106784992e-05, 1.3219284857637214e-05, 1.630553559996809e-05, 1.9825045445384375e-05, 2.3838636882425878e-05, 2.8415670881768902e-05, 3.3635245559400795e-05, 3.95875631124901e-05, 4.637548865068543e-05, 5.4116327861820987e-05, 6.294385423279893e-05, 7.301062085911147e-05, 8.44905967945913e-05, 9.758217350152745e-05, 0.00011251159335718234, 0.00012953685946652218, 0.0001489521943486774, 0.00017109312455002314, 0.00019634227905373608, 0.0002251360016911716, 0.0002579718918275978, 0.0002954174036368161, 0.00033811965257322685, 0.0003868165985121067, 0.00044234979881978883, 0.000505678951746213, 0.000577898481471814, 0.0006602564514236075, 0.00075417613271137, 0.000861280600419329, 0.0009834207828147287, 0.0011227074482063066, 0.0012815476822342878, 0.001462686485975001, 0.0016692542137387164, 0.0019048206703580484, 0.002173456802850449, 0.0024798050525806402, 0.0028291595837154453, 0.003227557774440908, 0.003681884552049176, 0.004199991374964317, 0.004790831917897655, 0.005464616804978631, 0.006232990064885314, 0.007109230357387858, 0.008108480448805902, 0.009248008902065284, 0.010547508503756575, 0.012029436585469348, 0.013719403120673117, 0.015646613304049737, 0.01784437226172206, 0.020350660614540016, 0.023208790841033847, 0.026468155782985957, 0.030185082228936854, 0.03442380432684787, 0.03925757364797396, 0.04476992508553923, 0.05105612046487287, 0.05822479481278606, 0.06639983373623022, 0.07572251435420017, 0.0863539467814424, 0.09847785835652628, 0.11230376872997776, 0.1280706106828306, 0.1460508592488553, 0.1665552404979834, 0.1899381013558711, 0.2166035332582641, 0.24701235546623276, 0.2816900787245926, 0.32123598688766164, 0.36633349345666244, 0.41776195200551264, 0.47641012459714155, 0.5432915409450616, 0.6195620137499113, 0.7065396129026261, 0.8057274437389025, 0.9188396229878619, 1.0478309013188452, 1.1949304444090645, 1.362680356320391, 1.5539796109278552, 1.7721341506018633, 2.020914017925626, 2.3046185077708192, 2.6281504656586163, 2.9971010163947933, 3.417846187218798, 3.897657095261897, 4.444825603521343, 5.068807616877284, 5.780386494525543, 6.591859402842075, 7.517249829140259, 8.572549928882209, 9.775996894473037]\n"
     ]
    }
   ],
   "source": [
    "print(radii)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
